"""
Feature Extraction and Visualization

提取并可视化模型中间层特征，支持：
- OverLoCK多层特征
- DINOv2 Transformer特征
- 自定义层Hook
"""
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')  # 非交互式后端
from typing import Dict, List, Optional, Tuple
import logging

logger = logging.getLogger(__name__)


class FeatureExtractor:
    """通用特征提取器"""

    def __init__(self, model: nn.Module, device: str = 'cuda'):
        """
        Args:
            model: PyTorch模型
            device: 运行设备
        """
        self.model = model
        self.device = device
        self.model.eval()

        # 存储特征的字典
        self.features: Dict[str, torch.Tensor] = {}
        self.hooks = []

    def register_hook(self, layer_name: str, module: nn.Module):
        """
        注册Hook到指定层

        Args:
            layer_name: 层名称（用于存储）
            module: 要Hook的模块
        """
        def hook_fn(module, input, output):
            # 保存输出特征
            self.features[layer_name] = output.detach()

        handle = module.register_forward_hook(hook_fn)
        self.hooks.append(handle)

    def clear_hooks(self):
        """清除所有Hook"""
        for hook in self.hooks:
            hook.remove()
        self.hooks = []
        self.features = {}

    def extract_features(self, input_tensor: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        提取特征

        Args:
            input_tensor: 输入张量 [B, C, H, W]

        Returns:
            特征字典 {layer_name: feature_tensor}
        """
        self.features = {}

        with torch.no_grad():
            _ = self.model(input_tensor)

        return self.features.copy()

    def visualize_feature_map(
        self,
        feature: torch.Tensor,
        num_channels: int = 32,
        nrow: int = 8,
        figsize: Tuple[int, int] = (16, 16)
    ) -> Image.Image:
        """
        可视化特征图

        Args:
            feature: 特征张量 [B, C, H, W]
            num_channels: 显示的通道数
            nrow: 每行显示的通道数
            figsize: 图像大小

        Returns:
            PIL.Image
        """
        # 取第一个batch
        if feature.dim() == 4:
            feature = feature[0]  # [C, H, W]

        # 限制通道数
        num_channels = min(num_channels, feature.shape[0])
        feature = feature[:num_channels]

        # 归一化到0-1
        feature = feature.cpu().numpy()
        feature_min = feature.min()
        feature_max = feature.max()
        if feature_max > feature_min:
            feature = (feature - feature_min) / (feature_max - feature_min)

        # 计算网格大小
        ncol = nrow
        nrow_actual = (num_channels + ncol - 1) // ncol

        # 创建图像
        fig, axes = plt.subplots(nrow_actual, ncol, figsize=figsize)
        if nrow_actual == 1 and ncol == 1:
            axes = np.array([[axes]])
        elif nrow_actual == 1:
            axes = axes.reshape(1, -1)
        elif ncol == 1:
            axes = axes.reshape(-1, 1)

        for idx in range(num_channels):
            row = idx // ncol
            col = idx % ncol
            ax = axes[row, col]

            ax.imshow(feature[idx], cmap='viridis')
            ax.set_title(f'Ch {idx}', fontsize=8)
            ax.axis('off')

        # 隐藏多余的子图
        for idx in range(num_channels, nrow_actual * ncol):
            row = idx // ncol
            col = idx % ncol
            axes[row, col].axis('off')

        plt.tight_layout()

        # 转换为PIL Image
        fig.canvas.draw()
        img_array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        img_array = img_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        img = Image.fromarray(img_array)

        plt.close(fig)

        return img


class OverLoCKFeatureExtractor(FeatureExtractor):
    """OverLoCK特征提取器"""

    def __init__(self, model, device='cuda'):
        super().__init__(model, device)

        # OverLoCK的关键层
        # 根据实际模型结构调整
        self.layer_names = {
            'stem': 'stem',
            'stage1': 'stages.0',
            'stage2': 'stages.1',
            'stage3': 'stages.2',
            'stage4': 'stages.3',
        }

    def setup_hooks(self):
        """设置OverLoCK的Hook"""
        self.clear_hooks()

        # 尝试注册各层Hook
        for layer_name, attr_path in self.layer_names.items():
            try:
                module = self._get_module_by_path(attr_path)
                if module is not None:
                    self.register_hook(layer_name, module)
                    logger.info(f"Registered hook for {layer_name}")
            except Exception as e:
                logger.warning(f"Failed to register hook for {layer_name}: {e}")

    def _get_module_by_path(self, path: str):
        """通过路径获取模块"""
        parts = path.split('.')
        module = self.model

        # 如果model有overlock_extractor属性
        if hasattr(self.model, 'overlock_extractor'):
            module = self.model.overlock_extractor

        for part in parts:
            if '[' in part:
                # 处理索引，如 stages[0]
                name, idx = part.split('[')
                idx = int(idx.rstrip(']'))
                module = getattr(module, name)[idx]
            else:
                module = getattr(module, part, None)
                if module is None:
                    return None

        return module


class DINOv2FeatureExtractor(FeatureExtractor):
    """DINOv2特征提取器"""

    def __init__(self, model, device='cuda'):
        super().__init__(model, device)

        # DINOv2关键层（Transformer blocks）
        self.layer_indices = [0, 3, 6, 9, 11]  # 选择几个代表性层

    def setup_hooks(self):
        """设置DINOv2的Hook"""
        self.clear_hooks()

        # 尝试获取DINO模型
        dino_model = None
        if hasattr(self.model, 'dino_extractor'):
            dino_model = self.model.dino_extractor
        elif hasattr(self.model, 'dinov2'):
            dino_model = self.model.dinov2

        if dino_model is None:
            logger.warning("DINOv2 model not found in refiner")
            return

        # 注册Transformer blocks的Hook
        if hasattr(dino_model, 'blocks'):
            for idx in self.layer_indices:
                if idx < len(dino_model.blocks):
                    layer_name = f'block_{idx}'
                    self.register_hook(layer_name, dino_model.blocks[idx])
                    logger.info(f"Registered hook for DINOv2 {layer_name}")


class FeatureVisualizer:
    """特征可视化统一接口"""

    def __init__(self, refiner_model, device='cuda'):
        """
        Args:
            refiner_model: Stage2 ROI Refiner模型
            device: 运行设备
        """
        self.refiner_model = refiner_model
        self.device = device

        # 创建提取器
        self.overlock_extractor = OverLoCKFeatureExtractor(refiner_model, device)
        self.dino_extractor = DINOv2FeatureExtractor(refiner_model, device)

    def visualize_overlock_features(
        self,
        image: torch.Tensor,
        roi_position: Optional[torch.Tensor] = None,
        num_channels: int = 32
    ) -> Dict[str, Image.Image]:
        """
        可视化OverLoCK特征

        Args:
            image: 输入图像 [1, 3, H, W]
            roi_position: ROI位置信息（可选）
            num_channels: 显示通道数

        Returns:
            {layer_name: PIL.Image}
        """
        # 设置Hook
        self.overlock_extractor.setup_hooks()

        # 提取特征
        if roi_position is not None:
            features = self.overlock_extractor.extract_features((image, roi_position))
        else:
            # 如果模型不需要roi_position
            try:
                features = self.overlock_extractor.extract_features(image)
            except:
                # 创建dummy roi_position
                roi_position = torch.tensor([[0.5, 0.5, 1.0, 1.0]], device=self.device)
                features = self.overlock_extractor.extract_features((image, roi_position))

        # 可视化每一层
        visualizations = {}
        for layer_name, feature in features.items():
            img = self.overlock_extractor.visualize_feature_map(
                feature,
                num_channels=num_channels,
                nrow=8
            )
            visualizations[layer_name] = img

        # 清除Hook
        self.overlock_extractor.clear_hooks()

        return visualizations

    def visualize_dino_features(
        self,
        image: torch.Tensor,
        roi_position: Optional[torch.Tensor] = None,
        num_channels: int = 32
    ) -> Dict[str, Image.Image]:
        """
        可视化DINOv2特征

        Args:
            image: 输入图像 [1, 3, H, W]
            roi_position: ROI位置信息（可选）
            num_channels: 显示通道数

        Returns:
            {layer_name: PIL.Image}
        """
        # 设置Hook
        self.dino_extractor.setup_hooks()

        # 提取特征
        if roi_position is not None:
            features = self.dino_extractor.extract_features((image, roi_position))
        else:
            try:
                features = self.dino_extractor.extract_features(image)
            except:
                roi_position = torch.tensor([[0.5, 0.5, 1.0, 1.0]], device=self.device)
                features = self.dino_extractor.extract_features((image, roi_position))

        # 可视化每一层
        visualizations = {}
        for layer_name, feature in features.items():
            # DINOv2输出可能是[B, N, D]格式，需要reshape
            if feature.dim() == 3:
                # [B, N, D] -> [B, D, sqrt(N), sqrt(N)]
                B, N, D = feature.shape
                H = W = int(np.sqrt(N))
                if H * W == N:
                    feature = feature.transpose(1, 2).reshape(B, D, H, W)
                else:
                    logger.warning(f"Cannot reshape {layer_name}: {feature.shape}")
                    continue

            img = self.dino_extractor.visualize_feature_map(
                feature,
                num_channels=min(num_channels, feature.shape[1]),
                nrow=8
            )
            visualizations[layer_name] = img

        # 清除Hook
        self.dino_extractor.clear_hooks()

        return visualizations


if __name__ == '__main__':
    # 测试代码
    logging.basicConfig(level=logging.INFO)

    print("Feature Extractor Module")
    print("This module provides feature extraction and visualization")
    print("Import and use in your code:")
    print("  from src.visualization.feature_extractor import FeatureVisualizer")
